import os
import re
import yaml
import json
import subprocess
import sys
import base64
from pathlib import Path
from flask import Flask, render_template, request, jsonify, send_from_directory, g
import librosa
from datetime import datetime

# 导入路径管理器和我们自己的模块
from src.utils import (
    CONFIG_PATH, DATABASE_PATH, MUSIC_LIBRARY_PATH, PROCESSOR_RESULTS_PATH,
    TEST_DATA_PATH, BATCH_RESULTS_PATH, PROJECT_ROOT, BATCH_HISTORY_PATH
)
from src.database import Database
from src.fingerprinting.processor import Processor
from src.fingerprinting.hashing import Hasher
from src.matching.scorer import Scorer
from src.tools.diagnostics import run_diagnostic_comparison
# [MERGED] 从新版本中加入 SampleGenerator 的顶层导入
from src.tools.sample_generator import SampleGenerator
from src.tools.diagnostics import run_diagnostic_comparison

# --- 初始化 ---
app = Flask(__name__, template_folder='web/templates', static_folder='web/static')
UPLOAD_FOLDER = Path('data/uploads')
TEST_DATA_PATH.mkdir(exist_ok=True)
UPLOAD_FOLDER.mkdir(exist_ok=True)
BATCH_HISTORY_PATH.mkdir(exist_ok=True)


with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
    config = yaml.safe_load(f)

processor_instance = Processor(config)
hasher_instance = Hasher(config['hashing'])
scorer_instance = Scorer(config['scoring'])

def get_db():
    """在应用上下文中管理数据库连接。"""
    if 'db' not in g:
        g.db = Database(str(DATABASE_PATH))
        g.db.connect()
    return g.db

@app.teardown_appcontext
def close_db(e=None):
    """在应用上下文结束时关闭数据库连接。"""
    db = g.pop('db', None)
    if db is not None:
        db.close()

# --- 页面与静态文件路由 ---

@app.route('/')
def index():
    """ 主页面 """
    return render_template('index.html')

@app.route('/plots/<path:filename>')
def serve_plot(filename):
    """提供可视化图表文件"""
    return send_from_directory(PROCESSOR_RESULTS_PATH, filename, as_attachment=False)

@app.route('/api/audio/<int:song_id>')
def serve_audio(song_id):
    """提供完整的歌曲音频文件"""
    db = get_db()
    song = db.get_song_by_id(song_id)
    if not song or not Path(song['file_path']).exists():
        return "Audio file not found.", 404
    audio_path = Path(song['file_path'])
    return send_from_directory(audio_path.parent, audio_path.name)

@app.route('/api/sample_audio/<path:filename>')
def serve_sample_audio(filename):
    """提供测试样本的音频文件。"""
    return send_from_directory(TEST_DATA_PATH, filename)


# --- 核心 API Endpoints ---

@app.route('/api/song_details/<int:song_id>')
def get_song_details(song_id):
    """获取单首歌曲的详细信息，并为其生成可视化图表。"""
    db = get_db()
    song = db.get_song_by_id(song_id)
    if not song: return jsonify({'error': 'Song not found'}), 404
    viz_data = processor_instance.process_and_visualize(song['file_path'])
    if not viz_data: return jsonify({'error': 'Failed to generate visualizations'}), 500
    return jsonify({**song, **viz_data})

def generate_candidate_visualization(song_path_str: str, offset: int, duration: float):
    """为识别出的候选匹配片段生成语谱图。"""
    processor_instance = Processor(config)
    try:
        song_path = Path(song_path_str)
        # 为每个片段创建唯一的输出目录
        output_dir = PROCESSOR_RESULTS_PATH / f"match_{song_path.stem}_{offset}s"
        output_dir.mkdir(exist_ok=True, parents=True)

        y, _ = librosa.load(song_path, sr=processor_instance.sr, mono=True)
        start_sample = int(offset * processor_instance.sr)
        end_sample = int((offset + duration) * processor_instance.sr)

        # 保证切片范围有效
        if start_sample < 0: start_sample = 0
        if end_sample > len(y): end_sample = len(y)
        if start_sample >= end_sample: return None

        y_slice = y[start_sample:end_sample]

        # 使用与主流程相同的处理步骤
        y_filtered = processor_instance._highpass_filter(y_slice)
        log_spec, freqs, times = processor_instance._compute_spectrogram(y_filtered)
        peaks = processor_instance._find_peaks(log_spec, freqs, times)

        spec_path = processor_instance._save_spectrogram_plot(
            log_spec, peaks,
            output_dir / "candidate_slice.png",
            f"匹配片段: {song_path.stem} at ~{offset}s"
        )
        return Path(spec_path).relative_to(PROCESSOR_RESULTS_PATH).as_posix()
    except Exception as e:
        print(f"Error generating candidate visualization: {e}")
        return None

@app.route('/api/recognize', methods=['POST'])
def recognize():
    """核心识别接口，处理上传的音频并返回匹配结果。"""
    filepath, sample_name = None, None
    db = get_db()
    try:
        if 'audio' in request.files:
            file = request.files['audio']
            if file.filename == '': return jsonify({'error': 'No selected file'}), 400
            sample_name = file.filename
            filepath = UPLOAD_FOLDER / sample_name
            file.save(filepath)
        elif request.json and 'filepath' in request.json:
            filepath_str = request.json['filepath']
            filepath = Path(filepath_str)
            sample_name = filepath.name
            if not filepath.exists(): return jsonify({'error': f'File not found on server at {filepath}'}), 404
        else:
            return jsonify({'error': 'No audio file or filepath provided'}), 400

        sample_duration = librosa.get_duration(path=filepath)
        sample_peaks = processor_instance.audio_to_peaks(str(filepath))
        sample_fingerprints = hasher_instance.peaks_to_fingerprints(sample_peaks)

        response_data = {
            'match': None, 'rejected_match': None, 'top_candidates': [], 'rejection_reason': None,
            'sample_analysis': {
                'name': sample_name,
                'peaks': sample_peaks.tolist() if sample_peaks.any() else [],
                'fingerprints_count': len(sample_fingerprints)
            }
        }

        if not sample_fingerprints: return jsonify(response_data)

        query_hashes = [fp[0] for fp in sample_fingerprints]
        db_matches = db.find_matches(query_hashes)
        result, diagnostics = scorer_instance.find_best_match(sample_fingerprints, db_matches)

        # 填充Top 5候选曲目信息
        histogram = diagnostics.get('histogram', {})

        if histogram:
            candidates = [{'song_id': sid, 'votes': sum(b.values()), 'offset': max(b, key=b.get)} for sid, b in histogram.items() if b]
            sorted_candidates = sorted(candidates, key=lambda x: x['votes'], reverse=True)[:5]
            for cand in sorted_candidates:
                song_info = db.get_song_by_id(cand['song_id'])
                if song_info:
                    cand['song_name'] = song_info.get('name', 'Unknown')
                    viz_url = generate_candidate_visualization(song_info['file_path'], cand['offset'], sample_duration)
                    cand['spectrogram_url'] = viz_url
            response_data['top_candidates'] = sorted_candidates

        # 填充最终匹配结果或被拒绝的匹配
        if result:
            song_info = db.get_song_by_id(result['song_id'])
            result['song_info'] = song_info
            response_data['match'] = result
        elif diagnostics.get('failed_report'):
            failed_report = diagnostics['failed_report']
            song_info = db.get_song_by_id(failed_report['song_id'])
            failed_report['song_info'] = song_info
            response_data['rejected_match'] = failed_report

        if diagnostics.get('rejection_reason'):
            response_data['rejection_reason'] = diagnostics['rejection_reason']

        return jsonify(response_data)

    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500
    finally:
        # 清理上传的临时文件
        if 'audio' in request.files and filepath and os.path.exists(filepath):
            os.remove(filepath)

# --- 其他工具API ---
@app.route('/api/stats')
def get_stats():
    """获取数据库统计信息。"""
    db = get_db()
    songs = db.get_all_songs()
    return jsonify({'track_count': len(songs)})

@app.route('/api/songs')
def get_songs():
    """获取所有歌曲列表。"""
    db = get_db()
    songs = db.get_all_songs()
    return jsonify(songs)

@app.route('/api/config/scoring', methods=['GET'])
def get_scoring_config():
    """获取当前的评分配置。"""
    return jsonify(config.get('scoring', {}))


# --- 工具类 API ---

@app.route('/api/tools/generate_samples', methods=['POST'])
def generate_samples():
    from src.tools.sample_generator import SampleGenerator
    generator = SampleGenerator(input_dir=MUSIC_LIBRARY_PATH, output_dir=TEST_DATA_PATH, config = config)
    generated = generator.generate_for_all()
    return jsonify({'message': f'成功生成 {len(generated)} 个测试样本。', 'files': generated})

@app.route('/api/tools/list_samples', methods=['GET'])
def list_samples():
    """列出现有的所有测试样本。"""
    if not TEST_DATA_PATH.exists(): return jsonify([])
    files = [f.name for f in TEST_DATA_PATH.glob('*.wav')]
    return jsonify(sorted(files, reverse=True))

@app.route('/api/tools/diagnose', methods=['POST'])
def diagnose_sample():
    """对单个样本进行诊断，与原曲进行深度对比。"""
    data = request.json
    sample_name = data.get('sample_name')
    if not sample_name: return jsonify({'error': 'Sample name not provided'}), 400

    match = re.match(r'(.+?)---(\d+)s---(\d+)s', sample_name.replace('.wav', ''))
    if not match: return jsonify({'error': '样本文件名格式不正确，无法解析时间信息。'}), 400

    original_stem, start_seconds_str, clip_length_str = match.groups()
    start_seconds = int(start_seconds_str)
    clip_length = int(clip_length_str)

    db = get_db()
    all_songs = db.get_all_songs()
    original_song = next((s for s in all_songs if Path(s['file_path']).stem == original_stem), None)

    if not original_song: return jsonify({'error': f"数据库中找不到名为 '{original_stem}' 的原曲。"}), 404

    original_path = Path(original_song['file_path'])
    sample_path = TEST_DATA_PATH / sample_name

    results = run_diagnostic_comparison(original_path, sample_path, start_seconds, clip_length, config)
    return jsonify(results)

@app.route('/api/tools/run_batch_test', methods=['POST'])
def run_batch_test_endpoint():
    """以后台进程方式启动批量测试脚本，支持传入自定义配置。"""
    try:
        custom_config = request.json
        script_path = str(PROJECT_ROOT / "scripts" / "05_batch_test.py")
        command = [sys.executable, script_path]

        # 如果有自定义配置，通过base64编码传递给子进程
        if custom_config and 'scoring' in custom_config:
            config_json = json.dumps(custom_config)
            config_b64 = base64.b64encode(config_json.encode('utf-8')).decode('utf-8')
            command.extend(['--custom-config-b64', config_b64])

        subprocess.Popen(command)
        return jsonify({'message': '批量测试已在后台启动。请稍后查看结果。'})
    except Exception as e:
        return jsonify({'error': '启动批量测试失败。', 'details': str(e)}), 500

@app.route('/api/batch_results/latest')
def get_latest_batch_results():
    """获取最新的批量测试结果。"""
    if not BATCH_RESULTS_PATH.exists(): return jsonify({'error': '尚未生成批量测试报告。'}), 404
    with open(BATCH_RESULTS_PATH, 'r', encoding='utf-8') as f:
        return jsonify(json.load(f))

@app.route('/api/batch_history')
def get_batch_history():
    """获取所有历史测试报告的列表。"""
    if not BATCH_HISTORY_PATH.exists(): return jsonify([])
    files_info = []
    for f in sorted(BATCH_HISTORY_PATH.glob('*.json'), reverse=True):
        try:
            timestamp_str = f.stem.replace('results_', '')
            dt_obj = datetime.strptime(timestamp_str, "%Y%m%d_%H%M%S")
            display_time = dt_obj.strftime("%Y-%m-%d %H:%M:%S")
        except ValueError:
            display_time = f.stem
        files_info.append({'filename': f.name, 'display_time': display_time})
    return jsonify(files_info)

@app.route('/api/batch_history/<string:filename>')
def get_specific_batch_result(filename):
    """根据文件名获取指定的历史测试报告。"""
    file_path = BATCH_HISTORY_PATH / filename
    if not file_path.exists(): return jsonify({'error': '指定的报告文件不存在。'}), 404
    with open(file_path, 'r', encoding='utf-8') as f:
        return jsonify(json.load(f))

if __name__ == '__main__':
    app.run(debug=True, port=6090)
